#lang racket

(require "util/utils.scm")
(require "definition-names.scm")
(require "scope-manager.scm")
(provide pass-scope)

;; inputs:
;;   the output language of pass-tag
;;   a list of tagged definitions
;;
;; outputs:
;;   lambdas are annotated with the free variables they capture
;;
;;   from: (lambda (<var> ...) <exp>)
;;                   ^ binders    ^ body
;;
;;   to:   (closure (<var> ...) ((var loc|env <var>) ...) (<var> ...) <exp>)
;;                    ^ caputures   ^ resolved captures      ^ binders  ^ body
;;
;;  furthermore all variable are resolved as one of
;;  * GLO - a global variable reference
;;  * LOC - a locally bound (stack) variable reference
;;  * TMP - a temporary (LET*) binding
;;  * ENV - an environment variable (closure) reference

(define (pass-scope source)
  (let ((globals (definition-names source)))
    (map* (lambda (def) (pass-scope-toplevel def globals)) source)))

(define (pass-scope-toplevel def globals)
  (match def
    (`(define ,name ,body) `(define ,name ,(pass-scope-exp body '() globals)))
    (`(dont-define ,body) `(dont-define ,(pass-scope-exp body '() globals)))
    (else (error 'pass-scope-toplevel "malformed toplevel definition" def))))

;; scope is a scope-manager
;;
(define (pass-scope-exp exp scope globals)
  (match exp
    (`(atomic ,x) `(atomic ,x))
    (`(var ,x)
     (cond ((member x globals) `(var glo ,x))
           ((scope-add-var! scope x) => (lambda (type) `(var ,type ,x)))
           (else `(var env ,x))))
    (`(app . ,exps) `(app . ,(mapply pass-scope-exp exps scope globals)))
    (`(begin . ,exps) `(begin . ,(mapply pass-scope-exp exps scope globals)))
    (`(lambda ,binds ,body)
     (let* ((scope^ (scope-create binds scope))
            (body (pass-scope-exp body scope^ (removes binds globals)))
            (captures (scope-list scope^))
            (resolved-captures (map (lambda (x) (pass-scope-exp `(var ,x) scope globals)) captures)))
       `(closure ,captures ,resolved-captures ,binds ,body)))
    (`(quote ,x) `(quote ,x))
    (`(if ,t ,c ,a)
     `(if ,(pass-scope-exp t scope globals)
          ,(pass-scope-exp c scope globals)
          ,(pass-scope-exp a scope globals)))
    (`(let* ,binds ,body)
     (let* ((bindings (let loop ((binds binds))
                        (if (null? binds)
                            '()
                            (let ((tmp (caar binds))
                                  (exp (pass-scope-exp (cadar binds) scope globals)))
                              (set! scope (scope-add-temporary tmp scope))
                              (cons (list tmp exp) (loop (cdr binds)))))))
            (body (pass-scope-exp body scope globals)))
       `(let* ,bindings ,body)))
    (`(letrec ,binds ,body)
     (begin
       (for-each (lambda (bind)
		   (set! scope (scope-add-temporary (car bind) scope)))
		 binds)
       (let* ((bindings (let loop ((binds binds))
			  (if (null? binds)
			      '()
			      (let ((tmp (caar binds))
				    (exp (pass-scope-exp (cadar binds) scope globals)))
				(cons (list tmp exp) (loop (cdr binds)))))))
	      (body (pass-scope-exp body scope globals)))
	 `(letrec ,bindings ,body))))
    (`(builtin ,nm . ,args)
     `(builtin ,nm . ,(mapply pass-scope-exp args scope globals)))
    (else (error 'pass-scope-exp "unknown type of tagged expression" exp))))
